"""
23. 合并 K 个升序链表
困难
相关标签
premium lock icon
相关企业
给你一个链表数组，每个链表都已经按升序排列。

请你将所有链表合并到一个升序链表中，返回合并后的链表。



示例 1：

输入：lists = [[1,4,5],[1,3,4],[2,6]]
输出：[1,1,2,3,4,4,5,6]
解释：链表数组如下：
[
  1->4->5,
  1->3->4,
  2->6
]
将它们合并到一个有序链表中得到。
1->1->2->3->4->4->5->6
示例 2：

输入：lists = []
输出：[]
示例 3：

输入：lists = [[]]
输出：[]


提示：

k == lists.length
0 <= k <= 10^4
0 <= lists[i].length <= 500
-10^4 <= lists[i][j] <= 10^4
lists[i] 按 升序 排列
lists[i].length 的总和不超过 10^4
"""
import heapq
from typing import Optional, List


# Definition for singly-linked list.
class ListNode:
    def __init__(self, val=0, next=None):
        self.val = val
        self.next = next
class Solution:
    def mergeKLists(self, lists: List[Optional[ListNode]]) -> Optional[ListNode]:
        l = []
        # 将每个链表的头节点加入堆中，格式：(值, 索引, 节点)
        for i, node in enumerate(lists):
            if node:
                heapq.heappush(l, (node.val, i, node))

        head = ListNode(0)
        cur = head
        while l:
            val,i,node = heapq.heappop(l)
            cur.next = node
            #取出下一节点
            nextNode = node.next
            # node.next = None
            cur = node
            if nextNode:
                heapq.heappush(l, (nextNode.val,i, nextNode))
        return head.next

    def mergeKLists_Res1(self, lists: List[Optional[ListNode]]) -> Optional[ListNode]:
        """
        最小堆，每次从list中取当前最小的结点，然后把当前最小结点的next加入到最小堆
        :param lists:
        :return:
        """
        # 创建最小堆
        heap = []
        # 将每个链表的头节点加入堆中
        for i, node in enumerate(lists):
            if node:
                heapq.heappush(heap, (node.val, i, node))

        # 创建虚拟头节点
        dummy = ListNode(0)
        current = dummy
        # 边遍历边添加
        while heap:
            val, i, node = heapq.heappop(heap)
            current.next = ListNode(val)
            current = current.next
            # 如果当前节点有下一个节点，则将其加入堆中
            if node.next:
                heapq.heappush(heap, (node.next.val, i, node.next))
        return dummy.next


def create_linked_list(arr):
    """辅助函数：根据数组创建链表"""
    if not arr:
        return None
    head = ListNode(arr[0])
    current = head
    for i in range(1, len(arr)):
        current.next = ListNode(arr[i])
        current = current.next
    return head


def linked_list_to_array(head):
    """辅助函数：将链表转换为数组便于验证"""
    result = []
    current = head
    while current:
        result.append(current.val)
        current = current.next
    return result


# 测试用例1：正常情况
def test_case_1():
    # 输入：lists = [[1,4,5],[1,3,4],[2,6]]
    list1 = create_linked_list([1, 4, 5])
    list2 = create_linked_list([1, 3, 4])
    list3 = create_linked_list([2, 6])
    lists = [list1, list2, list3]

    solution = Solution()
    result = solution.mergeKLists(lists)
    expected = [1, 1, 2, 3, 4, 4, 5, 6]
    print("test_case_1结果:",linked_list_to_array(result))
    print("test_case_1期望:",expected)
    assert linked_list_to_array(result) == expected
    print("测试用例1通过")


# 测试用例2：空列表
def test_case_2():
    # 输入：lists = []
    lists = []
    solution = Solution()
    result = solution.mergeKLists(lists)
    assert result is None
    print("测试用例2通过")


# 测试用例3：包含空链表
def test_case_3():
    # 输入：lists = [[]]
    lists = [None]
    solution = Solution()
    result = solution.mergeKLists(lists)
    assert result is None
    print("测试用例3通过")


# 测试用例4：单个链表
def test_case_4():
    # 输入：lists = [[1,2,3]]
    list1 = create_linked_list([1, 2, 3])
    lists = [list1]
    solution = Solution()
    result = solution.mergeKLists(lists)
    expected = [1, 2, 3]
    assert linked_list_to_array(result) == expected
    print("测试用例4通过")


# 测试用例5：多个空链表
def test_case_5():
    # 输入：lists = [[],[],[]]
    lists = [None, None, None]
    solution = Solution()
    result = solution.mergeKLists(lists)
    assert result is None
    print("测试用例5通过")


# 测试用例6：包含负数
def test_case_6():
    # 输入：lists = [[-1,0,1],[1,2,3],[-2,-1,0]]
    list1 = create_linked_list([-1, 0, 1])
    list2 = create_linked_list([1, 2, 3])
    list3 = create_linked_list([-2, -1, 0])
    lists = [list1, list2, list3]
    solution = Solution()
    result = solution.mergeKLists_Res1(lists)
    expected = [-2, -1, -1, 0, 0, 1, 1, 2, 3]
    assert linked_list_to_array(result) == expected
    print("测试用例6通过")


# 运行所有测试
if __name__ == "__main__":
    test_case_1()
    # test_case_2()
    # test_case_3()
    # test_case_4()
    # test_case_5()
    # test_case_6()
    print("所有测试用例通过！")



